import os
import shutil

import numpy as np
import matplotlib.pyplot as plt

'''
==================================================
kmeans的关键为如何表示样本之间的距离（核心问题）。
主要流程为:
1.手动设置簇的个数k。
2.在所有样本中随机选取k个样本作为簇的中心。
3.计算每个样本与簇中心的距离，并把该样本划分到最近的簇中。
4.更新簇的中心，计算每个簇的样本均值作为新的簇中心。
5.重复二、三步，直到簇的中心不再变化或变化很小。
==================================================
'''


def cas_iou(box,cluster):
    x = np.minimum(cluster[:,0],box[0])
    y = np.minimum(cluster[:,1],box[1])

    intersection = x * y
    area1 = box[0] * box[1]

    area2 = cluster[:,0] * cluster[:,1]
    iou = intersection / (area1 + area2 -intersection)

    return iou


def kmeans_anchors(boxes, k, ciou):
    '''
    Kmeans聚类算法，用于聚类yolov5的anchors

    Args:
        boxes: anchor的坐标数组，格式为(n, 2)，n为anchor数量，每个anchor有两个维度(width, height)
        k: 聚类簇数
        dist_func: 距离函数，本例中为1-Ciou距离函数

    Returns:
        聚类中心数组，格式为(k, 2)，每个聚类中心有两个维度(width, height)
    '''


    # 取出一共有多少框
    row = boxes.shape[0]

    # 初始化每个框各个点的位置
    distance = np.empty((row, k))

    # 最后的聚类位置
    last_clu = np.zeros((row,))

    # np.random.seed(0) # 设置随机数生成器的种子（seed），使得每次生成的随机数序列是可重复的。
    np.random.seed()

    # 随机选5个当聚类中心
    # np.random.choice(A, size=None, replace=True, p=None) 随机在A数组中抽取size个元素，replace为抽样完是否放回，P为抽取的概率
    cluster = boxes[np.random.choice(row, k, replace = False)]

    while True:
        # 计算每个样本距离哪个聚类中心最近
        '''
        np.apply_along_axis(func1d, axis, arr, *args, **kwargs)沿指定轴应用一个函数；
        · func1d函数只接受一维数组作为参数
        · axis是要在func1d函数上应用的轴
        · arr是要在func1d函数上应用的数组
        · *args, ** kwargs为func1d函数的参数
        
        函数的返回值是一个数组，其形状与arr相同，但是沿着指定的轴上的每个元素都被替换为该元素及其周围元素（沿着该轴）应用函数func1d的结果。
        '''


        # 计算每一行距离kg个簇中心的iou情况。
        for i in range(row):
            distance[i] = cas_iou(boxes[i], cluster)

        # 每个样本距离最近的簇
        near = np.argmin(distance, axis=1)


        if (last_clu == near).all():
            break

        # 求每一个类的中位点
        for j in range(k):
            cluster[j] = np.mean(boxes[near == j], axis=0)

        last_clu = near

    return cluster, near


def ciou_distance(box, cluster):
    """
    计算1-Ciou距离，用于聚类yolov5的anchors

    Args:
        box1: anchor1的坐标，格式为(2,)，两个维度为(width, height)
        box2: anchor2的坐标，格式为(2,)，两个维度为(width, height)

    Returns:
        1-Ciou距离
    """
    # w1, h1 = box1
    # w2, h2 = box2
    # area1, area2 = w1 * h1, w2 * h2

    # x = np.minimum(cluster[:,0],box[0])
    # y = np.minimum(cluster[:,1],box[1])
    #
    # intersection = x * y

    w1 = cluster[:,0]
    w2 = box[0]
    h1 = cluster[:,1]
    h2 = box[1]

    area1 = box[0] * box[1]
    area2 = cluster[:,0] * cluster[:,1]

    inter_area = np.minimum(cluster[:,0],box[0]) * np.minimum(cluster[:,1],box[1])
    union_area = area1 + area2 - inter_area
    iou = inter_area / union_area
    x1, y1 = w1 / 2, h1 / 2
    x2, y2 = w2 / 2, h2 / 2
    center_dist = (x1 - x2) ** 2 + (y1 - y2) ** 2
    # enclose_w, enclose_h = max(w1, w2), max(h1, h2)
    # enclose_area = enclose_w * enclose_h
    enclose_area = np.maximum(cluster[:,0],box[0]) * np.maximum(cluster[:,1],box[1])
    ciou = iou - (center_dist / enclose_area)
    return 1 - ciou


# def ciou_distance(boxes, clusters):
#     boxes = np.expand_dims(boxes, axis=1)
#     clusters = np.expand_dims(clusters, axis=0)
#
#     boxes_xy = boxes[..., :2]
#     boxes_wh = boxes[..., 2:4]
#     boxes_xy1 = boxes_xy - boxes_wh / 2
#     boxes_xy2 = boxes_xy + boxes_wh / 2
#
#     clusters_xy = clusters[..., :2]
#     clusters_wh = clusters[..., 2:4]
#     clusters_xy1 = clusters_xy - clusters_wh / 2
#     clusters_xy2 = clusters_xy + clusters_wh / 2
#
#     inter_xy1 = np.maximum(boxes_xy1, clusters_xy1)
#     inter_xy2 = np.minimum(boxes_xy2, clusters_xy2)
#     inter_wh = np.maximum(inter_xy2 - inter_xy1, 0)
#
#     inter_area = inter_wh[..., 0] * inter_wh[..., 1]
#     boxes_area = boxes_wh[..., 0] * boxes_wh[..., 1]
#     clusters_area = clusters_wh[..., 0] * clusters_wh[..., 1]
#
#     union_area = boxes_area + clusters_area - inter_area
#     iou = inter_area / union_area
#
#     center_distance = np.sum(np.square(boxes_xy - clusters_xy), axis=-1)
#     enclose_xy1 = np.minimum(boxes_xy1, clusters_xy1)
#     enclose_xy2 = np.maximum(boxes_xy2, clusters_xy2)
#     enclose_wh = np.maximum(enclose_xy2 - enclose_xy1, 0)
#
#     diagonal_distance = np.sum(np.square(enclose_wh), axis=-1)
#
#     ciou = iou - center_distance / diagonal_distance
#
#     return ciou


def lab_to_anchors(label_path, lab_total):
    txt_list = os.listdir(label_path)
    with open(lab_total, 'w') as f:
        for i in txt_list:
            if i.endswith('txt'):
                with open(label_path + '\\' + i, 'r', encoding = 'utf-8') as F:
                    F.seek(0)
                    txt_content = F.read()
                    print(txt_content)
                    f.write(str(txt_content))
                    pass
                pass
            pass
        pass


def pre_anchors(anchors):
    data = []
    array_temp = anchors[:]
    # a = len(array_temp[:])
    for i in range(len(array_temp[:])):
        # xmin = array_temp[i:i+1,[1]].flatten()
        # ymin = array_temp[i:i+1,[2]].flatten()
        # xmax = array_temp[i:i+1,[3]].flatten()
        # ymax = array_temp[i:i+1,[4]].flatten()
        # 得到宽高
        data.append([array_temp[-1], array_temp[-2]])

    return np.array(data)

if __name__ == '__main__':

    # 整合标签
    label_path = r'F:\yolov5_master_yyf\dataset\labels'
    lab_total = './lab_total.txt'
    # label_total = lab_to_anchors(label_path, lab_total) # 注意：这里需要手动创建lab_total.txt文件

    # 加载anchors
    anchors = np.loadtxt('./lab_total.txt')

    # anchors预处理  - (n, 5)->(n, 2)
    data = pre_anchors(anchors)
    data = data.reshape(-1, 2)

    # 聚类anchors
    k = 10
    centers, near = kmeans_anchors(data, k, ciou_distance)

    # 可视化聚类结果
    plt.scatter(anchors[:, 0], anchors[:, 1], s=1)
